"""
视频文件解析和处理模块
支持多种视频格式的解析和处理
"""

import os
import tempfile
from typing import Dict, Any, Optional, Tuple, List, Union
from pathlib import Path
import cv2
import numpy as np
from moviepy.editor import VideoFileClip, AudioFileClip, CompositeAudioClip
from dataclasses import dataclass
from enum import Enum
import librosa


class VideoFormat(Enum):
    """支持的视频格式"""
    MP4 = "mp4"
    AVI = "avi"
    MOV = "mov"
    WMV = "wmv"
    FLV = "flv"
    WEBM = "webm"
    MKV = "mkv"


@dataclass
class VideoInfo:
    """视频信息"""
    path: str
    duration: float
    fps: float
    width: int
    height: int
    channels: int
    audio_sample_rate: Optional[int] = None
    audio_channels: Optional[int] = None
    format: str = ""
    file_size: int = 0
    bitrate: Optional[int] = None
    codec: str = ""
    has_audio: bool = False


@dataclass
class AudioSegment:
    """音频片段"""
    start_time: float
    end_time: float
    text: str
    confidence: float = 1.0
    metadata: Dict[str, Any] = None


class VideoProcessor:
    """视频处理器"""
    
    def __init__(self):
        self.supported_formats = [
            '.mp4', '.avi', '.mov', '.wmv', '.flv', '.webm', '.mkv'
        ]
    
    def is_supported_format(self, file_path: Union[str, Path]) -> bool:
        """检查是否支持该视频格式"""
        return Path(file_path).suffix.lower() in self.supported_formats
    
    def get_video_info(self, video_path: Union[str, Path]) -> VideoInfo:
        """获取视频信息"""
        try:
            video_path = str(video_path)
            
            if not os.path.exists(video_path):
                raise FileNotFoundError(f"Video file not found: {video_path}")
            
            # 使用moviepy获取基本信息
            with VideoFileClip(video_path) as clip:
                duration = clip.duration
                fps = clip.fps
                width, height = clip.size
                
                # 音频信息
                has_audio = clip.audio is not None
                audio_sample_rate = None
                audio_channels = None
                
                if has_audio:
                    # 提取音频信息
                    audio_path = self._extract_audio_temp(clip)
                    if audio_path:
                        try:
                            y, sr = librosa.load(audio_path, sr=None)
                            audio_sample_rate = sr
                            audio_channels = 1 if len(y.shape) == 1 else y.shape[1]
                        finally:
                            if os.path.exists(audio_path):
                                os.unlink(audio_path)
                
                # 文件信息
                file_size = os.path.getsize(video_path)
                file_ext = Path(video_path).suffix.lower()
                
                return VideoInfo(
                    path=video_path,
                    duration=duration,
                    fps=fps,
                    width=width,
                    height=height,
                    channels=3,  # 假设RGB
                    audio_sample_rate=audio_sample_rate,
                    audio_channels=audio_channels,
                    format=file_ext.replace('.', ''),
                    file_size=file_size,
                    has_audio=has_audio
                )
                
        except Exception as e:
            raise RuntimeError(f"Failed to get video info: {str(e)}")
    
    def extract_audio(
        self, 
        video_path: Union[str, Path],
        output_path: Optional[Union[str, Path]] = None,
        start_time: Optional[float] = None,
        end_time: Optional[float] = None,
        sample_rate: int = 22050
    ) -> str:
        """提取视频中的音频"""
        try:
            video_path = str(video_path)
            
            # 生成输出路径
            if output_path is None:
                output_path = tempfile.mktemp(suffix='.wav')
            else:
                output_path = str(output_path)
                os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 使用moviepy提取音频
            with VideoFileClip(video_path) as clip:
                audio_clip = clip.audio
                
                if audio_clip is None:
                    raise ValueError("Video has no audio track")
                
                # 应用时间范围
                if start_time is not None or end_time is not None:
                    start = start_time or 0
                    end = end_time or audio_clip.duration
                    audio_clip = audio_clip.subclip(start, end)
                
                # 设置采样率
                if audio_clip.fps != sample_rate:
                    # 需要重新采样
                    temp_path = tempfile.mktemp(suffix='.wav')
                    audio_clip.write_audiofile(temp_path)
                    
                    # 使用librosa重新采样
                    y, sr = librosa.load(temp_path, sr=sample_rate)
                    import soundfile as sf
                    sf.write(output_path, y, sample_rate)
                    
                    # 清理临时文件
                    os.unlink(temp_path)
                else:
                    audio_clip.write_audiofile(output_path)
            
            return output_path
            
        except Exception as e:
            raise RuntimeError(f"Failed to extract audio: {str(e)}")
    
    def extract_frames(
        self,
        video_path: Union[str, Path],
        output_dir: Optional[Union[str, Path]] = None,
        frame_rate: Optional[float] = None,
        start_time: Optional[float] = None,
        end_time: Optional[float] = None
    ) -> List[str]:
        """提取视频帧"""
        try:
            video_path = str(video_path)
            
            # 生成输出目录
            if output_dir is None:
                output_dir = tempfile.mkdtemp()
            else:
                output_dir = str(output_dir)
                os.makedirs(output_dir, exist_ok=True)
            
            # 使用OpenCV提取帧
            cap = cv2.VideoCapture(video_path)
            
            if not cap.isOpened():
                raise RuntimeError("Failed to open video file")
            
            # 获取视频属性
            video_fps = cap.get(cv2.CAP_PROP_FPS)
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            duration = total_frames / video_fps
            
            # 计算时间范围
            start_frame = 0
            end_frame = total_frames
            
            if start_time is not None:
                start_frame = int(start_time * video_fps)
            if end_time is not None:
                end_frame = int(end_time * video_fps)
            
            # 计算帧间隔
            if frame_rate is None:
                frame_interval = 1
            else:
                frame_interval = int(video_fps / frame_rate)
            
            frame_paths = []
            frame_count = 0
            
            # 跳转到开始帧
            cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
            
            for i in range(start_frame, min(end_frame, total_frames), frame_interval):
                cap.set(cv2.CAP_PROP_POS_FRAMES, i)
                ret, frame = cap.read()
                
                if ret:
                    frame_path = os.path.join(output_dir, f"frame_{i:06d}.jpg")
                    cv2.imwrite(frame_path, frame)
                    frame_paths.append(frame_path)
                    frame_count += 1
            
            cap.release()
            return frame_paths
            
        except Exception as e:
            raise RuntimeError(f"Failed to extract frames: {str(e)}")
    
    def merge_audio_video(
        self,
        video_path: Union[str, Path],
        audio_path: Union[str, Path],
        output_path: Union[str, Path],
        start_time: float = 0.0,
        audio_volume: float = 1.0
    ) -> str:
        """合并音频和视频"""
        try:
            video_path = str(video_path)
            audio_path = str(audio_path)
            output_path = str(output_path)
            
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 加载视频和音频
            video_clip = VideoFileClip(video_path)
            audio_clip = AudioFileClip(audio_path)
            
            # 调整音频音量
            if audio_volume != 1.0:
                audio_clip = audio_clip.volumex(audio_volume)
            
            # 在指定时间插入音频
            if start_time > 0:
                # 创建静音段
                from moviepy.audio.fx.all import audio_fadeout
                silence = AudioFileClip(audio_path).subclip(0, 0)  # 创建空音频
                silence_with_duration = silence.set_duration(start_time)
                
                # 拼接音频
                final_audio = CompositeAudioClip([
                    silence_with_duration,
                    audio_clip.set_start(start_time)
                ])
            else:
                final_audio = audio_clip
            
            # 合并视频和音频
            final_clip = video_clip.set_audio(final_audio)
            
            # 如果音频比视频长，截断视频
            if final_audio.duration > video_clip.duration:
                final_clip = final_clip.subclip(0, video_clip.duration)
            elif final_audio.duration < video_clip.duration:
                # 如果音频比视频短，在音频结束后静音
                final_clip = final_clip
            
            # 导出最终视频
            final_clip.write_videofile(
                output_path,
                codec='libx264',
                audio_codec='aac',
                temp_audiofile='temp-audio.m4a',
                remove_temp=True
            )
            
            # 清理资源
            video_clip.close()
            audio_clip.close()
            final_clip.close()
            
            return output_path
            
        except Exception as e:
            raise RuntimeError(f"Failed to merge audio and video: {str(e)}")
    
    def replace_audio(
        self,
        video_path: Union[str, Path],
        new_audio_path: Union[str, Path],
        output_path: Union[str, Path]
    ) -> str:
        """替换视频的音频轨道"""
        try:
            video_path = str(video_path)
            new_audio_path = str(new_audio_path)
            output_path = str(output_path)
            
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 加载视频和新音频
            video_clip = VideoFileClip(video_path)
            new_audio_clip = AudioFileClip(new_audio_path)
            
            # 替换音频
            if new_audio_clip.duration < video_clip.duration:
                # 如果新音频较短，循环播放
                loops_needed = int(video_clip.duration / new_audio_clip.duration) + 1
                extended_audio = new_audio_clip.loop(loops_needed)
                final_audio = extended_audio.subclip(0, video_clip.duration)
            else:
                # 如果新音频较长，截断
                final_audio = new_audio_clip.subclip(0, video_clip.duration)
            
            # 设置新音频
            final_clip = video_clip.set_audio(final_audio)
            
            # 导出视频
            final_clip.write_videofile(
                output_path,
                codec='libx264',
                audio_codec='aac',
                temp_audiofile='temp-audio.m4a',
                remove_temp=True
            )
            
            # 清理资源
            video_clip.close()
            new_audio_clip.close()
            final_clip.close()
            
            return output_path
            
        except Exception as e:
            raise RuntimeError(f"Failed to replace audio: {str(e)}")
    
    def trim_video(
        self,
        video_path: Union[str, Path],
        output_path: Union[str, Path],
        start_time: float,
        end_time: float
    ) -> str:
        """裁剪视频"""
        try:
            video_path = str(video_path)
            output_path = str(output_path)
            
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 加载视频并裁剪
            with VideoFileClip(video_path) as clip:
                trimmed_clip = clip.subclip(start_time, end_time)
                trimmed_clip.write_videofile(output_path)
            
            return output_path
            
        except Exception as e:
            raise RuntimeError(f"Failed to trim video: {str(e)}")
    
    def concat_videos(
        self,
        video_paths: List[Union[str, Path]],
        output_path: Union[str, Path]
    ) -> str:
        """拼接视频"""
        try:
            output_path = str(output_path)
            
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 加载所有视频
            clips = []
            for video_path in video_paths:
                clip = VideoFileClip(str(video_path))
                clips.append(clip)
            
            # 拼接视频
            final_clip = clips[0]
            if len(clips) > 1:
                final_clip = clips[0].concatenate_videoclips(clips[1:])
            
            # 导出最终视频
            final_clip.write_videofile(output_path)
            
            # 清理资源
            for clip in clips:
                clip.close()
            
            return output_path
            
        except Exception as e:
            raise RuntimeError(f"Failed to concatenate videos: {str(e)}")
    
    def _extract_audio_temp(self, clip) -> Optional[str]:
        """提取音频到临时文件"""
        try:
            if clip.audio is None:
                return None
            
            temp_path = tempfile.mktemp(suffix='.wav')
            clip.audio.write_audiofile(temp_path)
            return temp_path
        except Exception:
            return None
    
    def get_video_thumbnail(
        self,
        video_path: Union[str, Path],
        output_path: Optional[Union[str, Path]] = None,
        timestamp: float = 1.0
    ) -> str:
        """生成视频缩略图"""
        try:
            video_path = str(video_path)
            
            if output_path is None:
                output_path = tempfile.mktemp(suffix='.jpg')
            else:
                output_path = str(output_path)
                os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 使用OpenCV生成缩略图
            cap = cv2.VideoCapture(video_path)
            
            if not cap.isOpened():
                raise RuntimeError("Failed to open video file")
            
            # 跳转到指定时间
            fps = cap.get(cv2.CAP_PROP_FPS)
            frame_number = int(timestamp * fps)
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
            
            ret, frame = cap.read()
            
            if ret:
                cv2.imwrite(output_path, frame)
            else:
                raise RuntimeError("Failed to capture frame")
            
            cap.release()
            return output_path
            
        except Exception as e:
            raise RuntimeError(f"Failed to generate thumbnail: {str(e)}")
    
    def validate_video_file(self, video_path: Union[str, Path]) -> bool:
        """验证视频文件是否有效"""
        try:
            video_path = str(video_path)
            
            if not os.path.exists(video_path):
                return False
            
            if not self.is_supported_format(video_path):
                return False
            
            # 尝试打开视频
            cap = cv2.VideoCapture(video_path)
            is_valid = cap.isOpened()
            cap.release()
            
            return is_valid
            
        except Exception:
            return False
